import torch
import torch.nn.functional as F
import matplotlib.pyplot as plt
import cv2
import os
from tqdm import tqdm
from concurrent.futures import ThreadPoolExecutor
from concurrent.futures import ProcessPoolExecutor
# 定义输入输出路径
test_file = '/root/autodl-tmp/Culane/list/test.txt'
output_folder = '/home/cc/NewPic'

# 读取文件内容
def read_tensor_data(file_path):
    with open(file_path, 'r') as file:
        data = file.read()
    tensor = torch.tensor
    tensor_data = eval(data.strip())
    return tensor_data

def draw_lines(image_path, gt_path, pred_path, output_path):
    # 读取图像
    image = cv2.imread(image_path)
    # image 剪裁到 [0, 1640, 270, 590]  x_min x_max y_min y_max
    # image = image[170:490, 0:1640]

    # 读取并绘制GT数据
    with open(gt_path, 'r') as f:
        gt_lines = f.readlines()
        for line in gt_lines:
            points = [max(0, int(float(x))) for x in line.strip().split()]
            for i in range(0, len(points) - 1, 2):
                if i + 3 < len(points):  # 确保有足够的点来画线
                    cv2.line(image, (points[i], points[i + 1]), (points[i + 2], points[i + 3]), (255, 0, 0), 2)  # 蓝色线

    # 读取并绘制预测数据
    with open(pred_path, 'r') as f:
        pred_lines = f.readlines()
        for line in pred_lines:
            points = [max(0, int(float(x))) for x in line.strip().split()]
            for i in range(0, len(points) - 1, 2):
                if i + 3 < len(points):  # 确保有足够的点来画线
                    cv2.line(image, (points[i], points[i + 1]), (points[i + 2], points[i + 3]), (0, 0, 255), 2)  # 红色线

    # 在图像上标注行数
    font = cv2.FONT_HERSHEY_SIMPLEX
    cv2.putText(image, f'GT Lines: {len(gt_lines)}', (10, 30), font, 1, (255, 0, 0), 2, cv2.LINE_AA)
    cv2.putText(image, f'Pred Lines: {len(pred_lines)}', (10, 70), font, 1, (0, 0, 255), 2, cv2.LINE_AA)

    return image

def process_line(line):
    line = line.strip()
    if not line.endswith('.jpg'):
        return

    image_path = '/root/autodl-tmp/Culane' + line
    base_name = os.path.splitext(line)[0]
    gt_path = '/root/autodl-tmp/Culane' + base_name + '.lines.txt'
    pred_path = '/home/cc/AAA_python_test/tmp&anchor/tmp_20250222_2113' + line.replace('.jpg', '.lines.txt')
    pred_anchor_path = '/home/cc/AAA_python_test/tmp&anchor/tmp_anchor_20250222_2113' + line.replace('.jpg', '.lines.txt')

    # 创建输出文件夹（如果不存在）
    output_dir = os.path.join(output_folder, os.path.dirname(line))
    os.makedirs(output_dir, exist_ok=True)

    # 处理图像并保存
    output_path = os.path.join(output_folder, line)
    image = draw_lines(image_path, gt_path, pred_path, output_path)

    # 读取tensor数据并绘制柱状图
    tensor_data = read_tensor_data(pred_anchor_path)

    # 获取组数
    num_groups = min(4, len(tensor_data[0]))

    if num_groups > 0:
        # 分割成三个长度为33的分布
        X = torch.stack([tensor_data[0][i][:33].clone().detach() for i in range(num_groups)])
        Y = torch.stack([tensor_data[0][i][33:66].clone().detach() for i in range(num_groups)])
        T = torch.stack([tensor_data[0][i][66:99].clone().detach() for i in range(num_groups)])

        # 应用softmax函数
        X_softmax = F.softmax(X, dim=1)
        Y_softmax = F.softmax(Y, dim=1)
        T_softmax = F.softmax(T, dim=1)

        # 绘制柱状图
        fig, axs = plt.subplots(3,num_groups, figsize=(10 * num_groups, 15))

        if num_groups == 1:
            for i in range(num_groups):
                axs[0].bar(range(33), X_softmax[i].numpy())
                axs[0].set_title(f'X Distribution - Group {i + 1}')

                axs[1].bar(range(33), Y_softmax[i].numpy())
                axs[1].set_title(f'Y Distribution - Group {i + 1}')

                axs[2].bar(range(33), T_softmax[i].numpy())
                axs[2].set_title(f'T Distribution - Group {i + 1}')
        else :
            for i in range(num_groups):
                axs[0, i].bar(range(33), X_softmax[i].numpy())
                axs[0, i].set_title(f'X Distribution - Group {i + 1}')

                axs[1, i].bar(range(33), Y_softmax[i].numpy())
                axs[1, i].set_title(f'Y Distribution - Group {i + 1}')

                axs[2, i].bar(range(33), T_softmax[i].numpy())
                axs[2, i].set_title(f'T Distribution - Group {i + 1}')

        plt.tight_layout()
        histogram_path = output_path.replace('.jpg', '_histogram.jpg')
        plt.savefig(histogram_path)
        plt.close()

        # 合并图像和柱状图
        histogram_image = cv2.imread(histogram_path)

        # 调整图像和柱状图的宽度一致
        if image.shape[1] != histogram_image.shape[1]:
            width = max(image.shape[1], histogram_image.shape[1])
            image = cv2.resize(image, (width, image.shape[0]))
            histogram_image = cv2.resize(histogram_image, (width, histogram_image.shape[0]))

        combined_image = cv2.vconcat([image, histogram_image])
        cv2.imwrite(output_path, combined_image)

        # 删除histogram.jpg
        os.remove(histogram_path)

# 读取 test.txt 文件
with open(test_file, 'r') as f:
    lines = f.readlines()

# 顺序处理每一行
# for line in tqdm(lines):
#     process_line(line)
#     print(line)
with ProcessPoolExecutor() as executor:
    list(tqdm(executor.map(process_line, lines), total=len(lines)))

print("处理完成！")
